import json
import asyncio
import websockets
import websockets.legacy.server as serve
import websockets.exceptions


class User:
    ws: serve.WebSocketServerProtocol
    userName: str  # 用户名
    opponentAddr: tuple  # 用于在字典中寻找对手
    # battleMap: list[list]  # 棋盘
    battleMap: list  # 棋盘

    def __init__(self, userName, ws):
        self.userName = userName
        self.ws = ws


def judgeWinStep(battleMap, row, col, rowStep, colStep) -> bool:
    """
    判断一个方向是否五子连珠
    """
    chess = battleMap[row][col]
    cnt = 1
    for i in range(1, mapSize):
        r = row - rowStep * i
        c = col - colStep * i
        if c not in range(mapSize) or r not in range(mapSize):
            break
        if battleMap[r][c] == chess:
            cnt += 1
        else:
            break

    for i in range(1, mapSize):
        r = row + rowStep * i
        c = col + colStep * i
        if c not in range(mapSize) or r not in range(mapSize):
            break
        if battleMap[r][c] == chess:
            cnt += 1
        else:
            break
    if cnt >= 5:
        return True
    else:
        return False


def judgeWin(battleMap, row, col) -> bool:
    """
    判断是否五子连珠
    只用判断最新的棋子周围即可
    """
    # 横向判断
    if judgeWinStep(battleMap, row, col, 0, 1):
        return True
    # 纵向判断
    if judgeWinStep(battleMap, row, col, 1, 0):
        return True
    # 斜向判断
    if judgeWinStep(battleMap, row, col, 1, 1):
        return True
    if judgeWinStep(battleMap, row, col, 1, -1):
        return True
    return False


async def battle(ws, chess: int):
    user = users[ws.remote_address]
    oppo = users[user.opponentAddr]
    print(ws.closed)
    print(ws.remote_address, users[ws.remote_address].opponentAddr)

    # 发送游戏初始化信息
    await ws.send(json.dumps({"type": gameStart, "content": oppo.userName, "state": chess == blackChess}))

    # 进入下棋循环
    while True:
        recv = json.loads(await ws.recv())
        if recv['type'] == move:
            # 接收新棋子信息
            row, col = recv['row'], recv['col']
            # 由于两个User的battleMap为浅引用 故只需修改一处
            user.battleMap[row][col] = chess
            print(user.userName, row, col)

            # 像双方分别发送棋子信息
            msg = json.dumps({"type": move, "row": row, "col": col, "color": chess})
            await oppo.ws.send(msg)
            await user.ws.send(msg)
            if judgeWin(user.battleMap, row, col):
                # 当前玩家胜利
                print(user.userName, ' win')
                # 发送游戏结束消息 并关闭连接
                msg = json.dumps({"type": gameEnd, "win": chess})
                await oppo.ws.send(msg)
                await user.ws.send(msg)
                await user.ws.close()
                await oppo.ws.close()
                users.pop(ws.remote_address)
                users.pop(user.opponentAddr)


async def login(ws: serve.WebSocketServerProtocol, path):
    global waitUser
    recv = await ws.recv()
    name = json.loads(recv)['content']
    print(recv)
    print(ws.remote_address)
    user = User(name, ws)
    users[ws.remote_address] = user

    try:
        await ws.send(json.dumps({"type": loginSuccess}))

        # 需要等待另一个用户对战
        if waitUser == ():
            waitUser = ws.remote_address
            user.battleMap = [[-1 for _ in range(mapSize)] for _ in range(mapSize)]
            async with waiting:
                await waiting.wait()
            print('awake')
            await battle(ws, blackChess)

        # 已经有用户等待
        else:
            async with waiting:
                waiting.notify(1)
            # 初始化自己的User对象
            user.opponentAddr = waitUser
            user.battleMap = users[waitUser].battleMap
            # 初始化对手的User对象
            users[waitUser].opponentAddr = ws.remote_address
            waitUser = ()
            await battle(ws, whiteChess)
    # 关闭连接后, 尝试通知对手
    except websockets.exceptions.ConnectionClosed:
        print(f"{name} closed")
        oppo = users.pop(user.opponentAddr, None)
        if oppo is not None and oppo.ws.open:
            await oppo.ws.send(json.dumps({"type": opponentQuit}))
    # 每个对战仅限一局, 结束后移除用户
    finally:
        users.pop(ws.remote_address, None)


def main():
    print('server start')
    start_server = websockets.serve(login, "0.0.0.0", 8000)

    asyncio.get_event_loop().run_until_complete(start_server)
    asyncio.get_event_loop().run_forever()


if __name__ == '__main__':
    mapSize = 10
    whiteChess = 0
    blackChess = 1
    loginSuccess = 1
    move = 2
    gameEnd = 3
    opponentQuit = 4
    gameStart = 5
    waitUser: tuple = ()
    waiting = asyncio.Condition()
    # users: dict[tuple, User] = {}
    users = {}
    main()
